package com.atguigu.aiproject.agent;

import com.atguigu.aiproject.service.ImgService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * 创作生成智能体
 * 用于根据用户输入自动生成推文等内容
 */
@Component
@Slf4j
public class TweetGenerationAgent {
    @Autowired
    @Qualifier("chatClient")
    private ChatClient chatClient;
    
    @Autowired
    private ImgService imgService;

    // 定义推文生成的提示模板
    private static final String TWEET_PROMPT_TEMPLATE = """
            你是一位专业的社交媒体内容创作者，擅长创作吸引人的推文。
            请根据以下主题创作一条推文：
            
            主题：{topic}
            
            要求：
            1. 推文长度在500个字符之间
            2. 内容要有吸引力，能够引起读者兴趣
            3. 语言风格要符合中文社交媒体习惯
            4. 可以适当使用表情符号增加表现力
            5. 如果主题不明确，请创作一条通用的有趣推文
            6. 内容应该包含有价值的信息或有趣的见解
            
            请直接输出推文内容，不要添加其他说明。
            """;

    // 定义图片提示词生成模板
    private static final String IMAGE_PROMPT_TEMPLATE = """
            根据以下推文内容，生成一个适合文中的图片描述：
            
            推文内容：{tweetContent}
            
            要求：
            1. 图片描述应与推文内容高度相关
            2. 描述要具体且生动，包含主体、场景、风格等要素
            3. 适合用于AI图像生成模型
            4. 不要包含文字或其他抽象概念
            5. 用中文描述
            
            请直接输出图片描述，不要添加其他说明。
            """;


    /**
     * 根据用户输入的主题生成推文（流式响应）
     *
     * @param topic 用户输入的主题
     * @return 推文内容的流式响应
     */
    public Flux<String> generateTweetStream(String topic) {
        PromptTemplate promptTemplate = new PromptTemplate(TWEET_PROMPT_TEMPLATE);
        Prompt prompt = promptTemplate.create(Map.of("topic", topic));

        return chatClient.prompt(prompt).stream().content();
    }
    /**
     * 根据推文内容生成图片提示词
     * 
     * @param tweetContent 推文内容
     * @return 图片提示词
     */
    public String generateImagePrompt(String tweetContent) {
        PromptTemplate promptTemplate = new PromptTemplate(IMAGE_PROMPT_TEMPLATE);
        Prompt prompt = promptTemplate.create(Map.of("tweetContent", tweetContent));
        return chatClient.prompt(prompt).call().content();
    }

    /**
     * 根据图片提示词生成单张图片
     * 
     * @param imagePrompt 图片提示词
     * @return 图片URL
     */
    public String generateImage(String imagePrompt) {
        // 调用ImgService生成图片
        return imgService.createSingleImage(imagePrompt);
    }
    
    /**
     * 根据图片提示词异步生成多张图片
     * 
     * @param imagePrompt 图片提示词
     * @param num 生成图片数量
     * @return 图片URL列表
     */
    public List<String> generateImages(String imagePrompt, int num) {
        // 调用ImgService异步生成多张图片，获取任务ID
        String taskId = imgService.createImages(imagePrompt, num);
        
        // 轮询获取图片状态，直到获取到图片或超时
        int maxAttempts = 10; // 最大尝试次数
        int attempt = 0;
        long delay = 2000; // 2秒延迟
        
        while (attempt < maxAttempts) {
            try {
                // 延迟一段时间再查询
                Thread.sleep(delay);
                
                // 根据任务ID获取图片URL列表
                List<String> imageUrls = imgService.getImageStatus(taskId);
                
                // 如果获取到图片URL，则返回
                if (imageUrls != null && !imageUrls.isEmpty()) {
                    return imageUrls;
                }
                
                attempt++;
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                break;
            } catch (Exception e) {
                log.warn("获取图片状态失败，尝试次数: {}", attempt + 1, e);
                attempt++;
            }
        }
        
        // 如果超时仍未获取到图片，返回空列表
        return new ArrayList<>();
    }


}